import * as vscode from 'vscode';
import * as fs from 'fs';
import * as path from 'path';

export interface ChatMessage {
    role: string;
    content: string;
    reasoning_content?: string;
}

export interface AssistantMessage extends ChatMessage {
    role: 'assistant';
    reasoning_content?: string;
}

export interface UserMessage extends ChatMessage {
    role: 'user';
}

export interface SystemMessage extends ChatMessage {
    role: 'system';
}

export interface ChatSession {
    id: string;
    title: string;
    createdAt: number;
    updatedAt: number;
    messages: (AssistantMessage | UserMessage | SystemMessage)[];
}

export class ChatHistoryManager {
    private context: vscode.ExtensionContext;
    private sessions: ChatSession[] = [];
    private currentSessionId: string | null = null;

    constructor(context: vscode.ExtensionContext) {
        this.context = context;
        this.loadSessions();
    }

    public async createNewSession(): Promise<ChatSession> {
        const session: ChatSession = {
            id: Date.now().toString(),
            title: '新会话',
            createdAt: Date.now(),
            updatedAt: Date.now(),
            messages: []
        };

        this.sessions.push(session);
        this.currentSessionId = session.id;
        await this.saveSessions();
        return session;
    }

    public async loadSessions(): Promise<ChatSession[]> {
        this.sessions = await this.context.globalState.get<ChatSession[]>('chatSessions', []);
        return this.sessions;
    }

    public async loadSession(sessionId: string): Promise<ChatSession | null> {
        const session = this.sessions.find(s => s.id === sessionId);
        if (session) {
            this.currentSessionId = sessionId;
            return session;
        }
        return null;
    }

    public async addMessage(message: AssistantMessage | UserMessage | SystemMessage): Promise<void> {
        const session = this.sessions.find(s => s.id === this.currentSessionId);
        if (session) {
            // 使用类型守卫来确保正确的类型推断
            const isAssistantMessage = (msg: any): msg is AssistantMessage => 
                msg.role === 'assistant';
            const isUserMessage = (msg: any): msg is UserMessage => 
                msg.role === 'user';
            const isSystemMessage = (msg: any): msg is SystemMessage => 
                msg.role === 'system';

            let typedMessage: AssistantMessage | UserMessage | SystemMessage;

            if (isAssistantMessage(message)) {
                typedMessage = {
                    role: 'assistant',
                    content: message.content,
                    reasoning_content: message.reasoning_content
                };
            } else if (isUserMessage(message)) {
                typedMessage = {
                    role: 'user',
                    content: message.content
                };
            } else if (isSystemMessage(message)) {
                typedMessage = {
                    role: 'system',
                    content: message.content
                };
            } else {
                throw new Error('Invalid message type');
            }

            session.messages.push(typedMessage);
            session.updatedAt = Date.now();
            
            // 如果是第一条用户消息，将其作为会话标题
            if (isUserMessage(message) && session.messages.length === 1) {
                session.title = message.content.slice(0, 20) + (message.content.length > 20 ? '...' : '');
            }
            
            await this.saveSessions();
        }
    }

    public async deleteSession(sessionId: string): Promise<void> {
        try {
            // 先加载最新的会话列表
            await this.loadSessions();
            
            const index = this.sessions.findIndex(s => s.id === sessionId);
            if (index !== -1) {
                // 删除会话
                this.sessions.splice(index, 1);
                
                // 如果删除的是当前会话，切换到最新的会话
                if (sessionId === this.currentSessionId) {
                    if (this.sessions.length > 0) {
                        this.currentSessionId = this.sessions[this.sessions.length - 1].id;
                    } else {
                        // 如果没有会话了，创建新会话
                        const newSession = await this.createNewSession();
                        this.currentSessionId = newSession.id;
                    }
                }
                
                // 立即保存更新后的会话列表到 VSCode 存储
                await this.context.globalState.update('chatSessions', this.sessions);
                
                // 确保更改已保存
                await this.saveSessions();
            }
        } catch (error) {
            console.error('Error deleting session:', error);
            throw error;
        }
    }

    public async renameSession(sessionId: string, title: string): Promise<void> {
        const session = this.sessions.find(s => s.id === sessionId);
        if (session) {
            session.title = title;
            await this.saveSessions();
        }
    }

    private async saveSessions(): Promise<void> {
        try {
            await this.context.globalState.update('chatSessions', this.sessions);
        } catch (error) {
            console.error('Error saving sessions:', error);
            throw error;
        }
    }

    public getCurrentSessionId(): string | null {
        return this.currentSessionId;
    }

    public getCurrentSession(): ChatSession | null {
        return this.sessions.find(s => s.id === this.currentSessionId) || null;
    }
} 